#!/usr/bin/env python
import argparse
import os
import sys
import re
import shutil
class inpGen():
    def __init__(self):
        self.title=''
        self.directory=''
     #accuracy  
        self.NBANDS=''
        self.ENCUT=''
        self.KPOINTS='2,1,1'
        self.PRECFOCK='Fast'
        self.PREC='Normal'
        self.LREAL='.FALSE'
      #accuracy  gw   
        self.NOMEGA='20'
        self.ENCUTGW='100'
        self.NBANDSGW=''
        self.NBANDSGWLOW=''
    #switch 
        self.LORBIT='10'
        self.LOPTICS='.FALSE'
        self.LWANNIER90=''
        self.NPAR=''
        self.NPROC=16
        self.SMEARSIGMA='-5,0.2'
        self.CUSTOM=''
        self.istart=''
    #JOB TYPES
        self.SCF='ALGO = Damped;TIME = 0.4'
        self.HSE03='LHFCALC=.TRUE;AEXX=0.25;HFSCREEN=0.3'
        self.HSE06='LHFCALC=.TRUE;AEXX=0.25;HFSCREEN=0.2'
        self.PBE0='LHFCALC=.TRUE;AEXX=0.25'
        self.G0W0='ALGO=GW0'
        self.BSE='ALGO=BSE'
    #optimization
        self.NSW=''
        self.ISIF=''
        self.IBRION=''
        self.POTIM=''
    #kpath
        self.nkpoints=100
        self.kpathlist=[]
        self.customK=0
        
    def writeKPOINTS(self,number):
        kpoints=open('KPOINTS_'+str(number),'w')       
        kpoints.write(self.title+'\n')
        kpoints.write('0\n')
        kpoints.write('Gamma\n')
        kpoints.write(' '.join(self.KPOINTS.split(',')))
        kpoints.close()  
        
    def readIBZKPT(self,number):
        try:
            ibzkpt=open('IBZKPT','r')
        except IOError:
            print 'IBZKPT file not found, can not do BAND calculation'
            sys.exit()
        kpoints=open('KPOINTS_'+str(number),'w')
        count=0
        pline=''
        cline=''
        addKpoints=[]
        if self.kpathlist:
            print self.kpathlist
            num_customK=len(self.kpathlist)
            spanK=round(self.nkpoints/(num_customK-1))
            self.customK=spanK*(num_customK-1)+1

        for line in ibzkpt:
            if re.search('^\s+[0-9]+\s*$',line):
                originK=float(line.split()[0])
                if self.customK:
                    acturalK=originK+self.customK
                else:
                    spanK=round(self.nkpoints/(originK-1))
                    acturalK=spanK*(originK-1)+1+originK
                line='  '+str(int(acturalK))+'\n'

            if re.search('[0-9]\.[0-9]+\s+[0-9]\.[0-9]+\s+[0-9]\.[0-9]+\s+',line) and not self.customK:
                self.kpathlist.append(line)
                ###append kpath to
            kpoints.write(line)
             
        for line in self.kpathlist:  
            if pline=='':
                pline=line
                addKpoints.append(str(float(pline.split()[0]))+'  '+str(float(pline.split()[1]))+'  '+str(float(pline.split()[2]))+'   0\n')
            else:
                cline=line
                clineX=float(cline.split()[0])
                clineY=float(cline.split()[1])
                clineZ=float(cline.split()[2])
                plineX=float(pline.split()[0])
                plineY=float(pline.split()[1])
                plineZ=float(pline.split()[2])
                for j in range(int(spanK)):
                    x=(clineX-plineX)/spanK*(j+1)+plineX
                    y=(clineY-plineY)/spanK*(j+1)+plineY
                    z=(clineZ-plineZ)/spanK*(j+1)+plineZ
                    addKpoints.append(str(x)+'  '+str(y)+'  '+str(z)+'   0\n')
                pline=cline  
        for kp in addKpoints:
            kpoints.write(kp)
            
    def writeINCAR(self,number,jobType='PBE',):
        if jobType in ['BAND']:
            self.readIBZKPT(number)
            try : shutil.copy('INCAR_'+str(number-1),'INCAR_'+str(number))
            except IOError:
                shutil.copy('INCAR','INCAR_'+str(number))
            incar=open('INCAR_'+str(number),'a')
            incar.write('ICHARGE=11')
            incar.close()
        else:
            incar=open('INCAR_'+str(number),'w')
            incar.write('SYSTEM='+self.title+'\n')
            incar.write('LORBITAL='+self.LORBIT+'\n')
            
            if self.PREC:
                incar.write('PREC='+self.PREC+'\n')
            if self.ENCUT:
                incar.write('ENCUT='+self.ENCUT+'\n')
            if self.istart:
                incar.write('ISTART='+self.istart+'\n')
            if self.SMEARSIGMA.split(',')[0]:
                incar.write('ISMEAR='+self.SMEARSIGMA.split(',')[0]+'\n')
            if self.SMEARSIGMA.split(',')[1]:
                incar.write('SIGMA='+self.SMEARSIGMA.split(',')[1]+'\n')
            if self.CUSTOM:
                incar.write(self.CUSTOM+'\n')
            if self.NBANDS:
                incar.write('NBANDS='+self.NBANDS+'\n')
            if self.LWANNIER90:
                incar.write('LWANNIER90='+self.LWANNIER90+'\n')
        if jobType in ['PBE']:
            incar.write('LREAL=AUTO\n')
        if jobType in ['HSE06','HSE03','PBE0']:
            incar.write('#Hybride functional input\n')
            incar.write(self.SCF+'\n')
            incar.write('PRECFOCK='+self.PRECFOCK+'\n')
            incar.write('LOPTICS='+self.LOPTICS+'\n') 
            incar.write(eval('self.'+jobType)+'\n')  
        if jobType=='G0W0':
            incar.write('#GW input\n')
            incar.write(self.G0W0+'\n')
            incar.write('NOMEGA='+self.NOMEGA+'\n')
            incar.write('ENCUTGW='+self.ENCUTGW+'\n')
            incar.write('PRECFOCK='+self.PRECFOCK+'\n')
            if self.NBANDSGW:   
                incar.write('NBANDSGW='+self.NBANDSGW+'\n')
            if self.NBANDSGWLOW:
                incar.write('NBANDSGWLOW='+self.NBANDSGWLOW+'\n')
        else:
            if self.NSW:
                incar.write('NSW='+self.NSW+'\n')
            if self.IBRION:
                incar.write('IBRION='+self.IBRION+'\n')
            if self.ISIF:
                incar.write('ISIF='+self.ISIF+'\n')
parser=argparse.ArgumentParser(formatter_class=argparse.RawTextHelpFormatter,description="A script used to generate vasp input")            
parser.add_argument("-j","--job",help="Set the job type (PBE,HSE03,HSE06,PBE0,G0W0,BSE,BAND), seperate multiple job by \";\". one job correspond to one INCAR file",  type=str,default='')
parser.add_argument("-k","--kpoints",help="Set the kpoints. a three number seperate by comma, e.g. 3,1,1;6,1,1.",  type=str,default='')  
parser.add_argument("-N","--nkpoints", help="set the number of total K points in BAND structure calculation, default:100",  type=int,default=100)
parser.add_argument("-K ","--kpointPath", help="Set the custom k point path in BAND structure calculation. format:\"0 0 0;0.5 0 0;0.5 0.5 0\" default path is read from IBZKPT file", type=str) 
parser.add_argument("-e","--encut",help="set the energy cut off. Default: None", type=str,default='')
parser.add_argument("-b","--bands",help="Set the number of bands. Default: do not set",  type=str,default='') 
parser.add_argument("-p","--prec",help="Set PREC of the jobs(Normal,Single,Accurate,Low,Med,High).Default:Normal", type=str,default='Normal') 
parser.add_argument("-i","--istart",help="Set istart value:(0 new job,1 restart,2),default:None",type=str,default='') 
parser.add_argument("-s","--smearsigma",help="Set the smear method and sigma, seperate by comma. Default: 0,0.05",  type=str,default='0,0.05') 
parser.add_argument("-W","--wannier",help="turn on out put option for wannier, anything:on, nothing:off", type=str,default='')
parser.add_argument("-o","--isif",help="Set the ISIF value, default:None. Set this option will also set ibrion to 1 and nsw to 100.(1,2:atoms only; 3:everything; 4:constrain volume; 5:cell only, fix volume; 6:cell only; 7:volume only)" ,type=str,default='')
parser.add_argument("-c","--custom",help="Set custom keywords. Seperate multiple keywords by comma. e.g.\"PREFOCK=Fast,ENCUT=200\"", type=str,default='')  
parser.add_argument("-n","--nproc",help="Set the number of cores to run the jobs, Default:16", type=str,default='16') 
parser.add_argument("-E","--encutgw",help="set the energy cut off for gw. Default: 100", type=str,default='200')
parser.add_argument("-w","--nomega",help="Set the number of omega for gw calculation, Default: 20",  type=str,default='20') 
parser.add_argument("--bandsgw",help="Set the number of bands for gw calculation. Default: not set", type=str,default='') 
parser.add_argument("--nbandsgwlow",help="Set the lowest band included in GW, Default:None", type=str,default='')  
parser.add_argument("--ibrion",help="Set ibrion value:(0:dynamics; 1:quasi-Newton; 2:conjugate-gradient; 3:damped dynamics; 5,6:numerical freq; 7,8:DFPT freq),default:None",type=str,default='') 
parser.add_argument("--nsw",help="Set number of ionic steps,default:None",type=str,default='') 
parser.add_argument("targetDIR",nargs='*')         
args=parser.parse_args()

#parse parameters
jobs=args.job.upper().split(';')
step=len(jobs)
kpoints=args.kpoints.split(';')
kpoints=kpoints+(step-len(kpoints))*['']
bands=args.bands.split(';')
bands=bands+(step-len(bands))*['']
encutgw=args.encutgw.split(';')
encutgw=encutgw+(step-len(encutgw))*['']
nomega=args.nomega.split(';')
nomega=nomega+(step-len(nomega))*['']
bandsgw=args.bandsgw.split(';')
bandsgw=bandsgw+(step-len(bandsgw))*['']
smearsigma=args.smearsigma.split(';')
smearsigma=smearsigma+(step-len(smearsigma))*['']
wannier=args.wannier.split(';')
wannier=wannier+(step-len(wannier))*['']
prec=args.prec.split(';')
prec=prec+(step-len(prec))*['']
encut=args.encut.split(';')
encut=encut+(step-len(encut))*['']
custom=args.custom.split(';')
custom=custom+(step-len(custom))*['']
isif=(step-1)*['']+list(args.isif)
if not args.isif:
    isif=['']*len(jobs)
    #if isif is set to 1 or 2 or 3 set nsw and ibrion
ibrion=''
nsw=''
if args.isif in list('123'):
    ibrion=list('1')
    nsw=['100']
if args.ibrion:
    ibrion=list(args.ibrion)
elif not ibrion:
    ibrion=['']
ibrion=(step-1)*['']+ibrion
if args.nsw:
    nsw=args.nsw.split('   ')
elif not nsw:
    nsw=['']
nsw=(step-1)*['']+nsw



incar=inpGen()
incar.nkpoints=args.nkpoints
if args.kpointPath:
    incar.kpathlist=args.kpointPath.split(';')

if 'G0W0' in jobs:
    incar.LOPTICS='.TRUE'

if args.targetDIR:
    orginDIR=os.getcwd()
    for targetDIR in args.targetDIR:
        if not os.path.exists(targetDIR):
            print "Error!!! Target Dierctory "+targetDIR+" do not exist!"
            sys.exit()
        os.chdir(targetDIR)   
        runvasp=open('runvasp.sh','w')
        for i in range(len(jobs)):
            if kpoints[i]:
                incar.KPOINTS=kpoints[i]
            if bands[i]:
                incar.NBANDS=bands[i]
            if encutgw[i]:
                incar.ENCUTGW=encutgw[i]
            if nomega[i]:
                incar.NOMEGA=nomega[i]
            if bandsgw[i]:
                incar.NBANDSGW=bandsgw[i]
            if wannier[i]:
                incar.LWANNIER90='.True'
            else:
                incar.LWANNIER90=''
            if custom[i]:
                incar.CUSTOM=custom[i].replace(',',';')
            if smearsigma[i]:
                incar.SMEARSIGMA=smearsigma[i]
            if prec[i]:
                incar.PREC=prec[i]
            if encut[i]:
                incar.ENCUT=encut[i]
            if i==0 and args.istart:
                incar.istart=args.istart
            else:
                incar.istart=''   
            if isif[i]:
                incar.ISIF=isif[i]
            if ibrion[i]:
                incar.IBRION=ibrion[i]
            if nsw[i]:
                incar.NSW=nsw[i]
            incar.writeKPOINTS(i)
            incar.writeINCAR(i,jobs[i])
            runvasp.write('cp INCAR_'+str(i)+' INCAR\n')
            runvasp.write('cp KPOINTS_'+str(i)+' KPOINTS\n')
            runvasp.write('mpirun -np '+args.nproc+' vasp\n')
        os.chdir(orginDIR)
else:
    runvasp=open('runvasp.sh','w')
    for i in range(len(jobs)):
        if kpoints[i]:
            incar.KPOINTS=kpoints[i]
        if bands[i]:
            incar.NBANDS=bands[i]
        if encutgw[i]:
            incar.ENCUTGW=encutgw[i]
        if nomega[i]:
            incar.NOMEGA=nomega[i]
        if bandsgw[i]:
            incar.NBANDSGW=bandsgw[i]
        if wannier[i]:
            incar.LWANNIER90='.True'
        else:
            incar.LWANNIER90=''
        if custom[i]:
            incar.CUSTOM=custom[i].replace(',',';')
        if smearsigma[i]:
            incar.SMEARSIGMA=smearsigma[i]
        if prec[i]:
            incar.PREC=prec[i]
        if encut[i]:
            incar.ENCUT=encut[i]    
        if i==0 and args.istart:
            incar.istart=args.istart 
        else:
            incar.istart='' 
        #print i  
        #print isif
        if isif[i]:
            incar.ISIF=isif[i]
        if ibrion[i]:
            incar.IBRION=ibrion[i]
        if nsw[i]:
            incar.NSW=nsw[i]
        incar.writeKPOINTS(i)
        incar.writeINCAR(i,jobs[i])
        runvasp.write('cp INCAR_'+str(i)+' INCAR\n')
        runvasp.write('cp KPOINTS_'+str(i)+' KPOINTS\n')
        runvasp.write('mpirun -np '+args.nproc+' vasp\n')

#
#test.writeINCAR(1)         
            
            
        
        
        
        
